import torch
from torchvision import datasets
import matplotlib.pyplot as plt

import sys
sys.path.insert(0,'F:\lab\\adv_mnist\DeepRobust-master') # 将该路径插入到包的搜索路径中
print(sys.path)

import deeprobust.image.netmodels.resnet as resnet


model = resnet.ResNet18([1,28,28]).to('cuda')

model.load_state_dict(torch.load("./trained_models/MNIST/MNIST_ResNet18_epoch_20.pt"))
model.eval()

xx = datasets.MNIST('./dataSet/MNIST', download = True).data[5:6].to('cuda')
xx = xx.unsqueeze_(1).float()/255
#print(xx.size())

## Set Target
yy = datasets.MNIST('./dataSet/MNIST', download = True).targets[5:6].to('cuda')

predict0 = model(xx)
predict0= predict0.argmax(dim=1, keepdim=True)

xx = xx.cpu().detach().numpy()  #gpu上的数据复制到cpu
yy = yy.cpu().detach().numpy()
print("yy label:{}".format(yy))

print("original prediction:")
print(predict0)

plt.imshow(xx[0,0]*255,cmap='gray',vmin=0,vmax=255) #使用data生成图片
plt.show()  #显示图片